# radiosity_relaxation.py
#
# Functions to apply relaxation techniques to solving the radiosity equation

from xmlscene_functions import *
from triangulatedpolygon import TriangulatedPolygon
from patch import *
from transformations import *
from gl_functions import *

from formfactor import *

import unittest
from copy import copy

class RadiosityRelaxationNotifier():
    def __init__(self, reporting_interval=1):
        self.guess_updates = {}
        self.guess_updates_name = "guess_updates"
        self.residual_updates = {}
        self.residual_updates_name = "residual_updates"
        self.convergence_metric_updates = {}
        self.convergence_metric_updates_name = "convergence_metric_updates"
        self.reporting_interval = reporting_interval
        self.intervals = {}

    def shouldReport(self, dictionary_name, relaxation_function):
        interval = 0
        if self.intervals.has_key(relaxation_function):
            if self.intervals[relaxation_function].has_key(dictionary_name):
                interval = self.intervals[relaxation_function][dictionary_name]
            else:
                self.intervals[relaxation_function][dictionary_name] = 1
                interval = 1
        else:
            self.intervals[relaxation_function] = {}
            self.intervals[relaxation_function][dictionary_name] = 1
            interval = 1

        report = (interval >= self.reporting_interval)
        if report:
            self.intervals[relaxation_function][dictionary_name] = 1
        else:
            self.intervals[relaxation_function][dictionary_name] += 1

        return report

    def appendValueToKeyInDictionary(self, dictionary, key, value):
        if dictionary.has_key(key):
            dictionary[key].append([value, len(dictionary[key])])
        else:
            dictionary[key] = []
            dictionary[key].append([value, 1])
    
    def notifyUpdatedGuess(self, guess, relaxation_function):
        if self.shouldReport(self.guess_updates_name, relaxation_function):
            saved_guess = copy(guess)
            self.appendValueToKeyInDictionary(self.guess_updates, relaxation_function, saved_guess)
        
    def notifyUpdatedResidual(self, residual, relaxation_function):
        if self.shouldReport(self.residual_updates_name, relaxation_function):
            saved_residual = copy(residual)
            self.appendValueToKeyInDictionary(self.residual_updates, relaxation_function, saved_residual)

    def notifyUpdatedConvergenceMetric(self, convergence, relaxation_function):
        if self.shouldReport(self.convergence_metric_updates_name, relaxation_function):
            saved_convergence = copy(convergence)
            self.appendValueToKeyInDictionary(self.convergence_metric_updates, relaxation_function, saved_convergence)

def radiosityRelaxation_JacobiIteration(patches, formfactorengine, initial_guess=None, convergence_threshold="0.00001,0.00001,0.00001", cached_form_factors=None, formfactor_threshold=0.0, notifier=None):
    guess = initial_guess
    if guess is None:
        guess = {}
        for patch in patches:
            guess[patch] = patch.get_emission()
    new_guess = copy(guess)

    # during each step, we compute a new guess using the previous radiosity values,
    # and then carry those new values into the next iteration
    len_patches = len(patches)
    converged = False
    while not converged:
        # compute the current step's guess
        for i in xrange(len_patches):
            patch = patches[i]

            # compute form factors
            if cached_form_factors.has_key(patch):
                form_factors = cached_form_factors[patch]
            else:
                form_factors = formfactorengine.computeFormFactorsForPatch(patch)

            # compute K_ii (self form-factor)
            if form_factors.has_key(patch):
                F_ii = form_factors[patch]
            else:
                F_ii = formfactorengine.defaultFormFactor()
            K_ii = subtractRGBStrings(whiteRGBString(), multiplyRGBStringBy(patch.get_reflectivity(), F_ii))

            # B^(k+1) = E - K_j/K_ii*B_j^k
            new_guess[patch] = patch.get_emission()
            rho_i = patch.get_reflectivity()
            for patch_index in form_factors:
                if patch_index == i or patch_index >= len_patches:
                    continue

                # compute K_ji
                F_ij = form_factors[patch_index]
                K_ij = multiplyRGBStringBy(rho_i, -F_ij)
                
                new_guess[patch] = subtractRGBStrings(new_guess[patch], multiplyRGBStrings(guess[patches[patch_index]], divideRGBStrings(K_ij, K_ii)))

        # save the current guess for the next step's calculations
        max_soln_change = blackRGBString()
        for patch in guess:
            # track change for convergence testing
            change = subtractRGBStrings(new_guess[patch], guess[patch])
            if isRGBStringGreaterThan(change, max_soln_change):
                max_soln_change = change

            # update the new radiosity value
            guess[patch] = new_guess[patch]

        # test for convergence
        if isRGBStringGreaterThan(convergence_threshold, max_soln_change):
            converged = True

        # notify after all the patches have been updated in this step
        if notifier is not None:
            notifier.notifyUpdatedGuess(guess, radiosityRelaxation_JacobiIteration)
            notifier.notifyUpdatedConvergenceMetric(max_soln_change, radiosityRelaxation_JacobiIteration)

    return guess

def radiosityRelaxation_GaussSeidelIteration(patches, formfactorengine, initial_guess=None, convergence_threshold="0.00001,0.00001,0.00001", cached_form_factors=None, formfactor_threshold=0.0, notifier=None):
    guess = initial_guess
    if guess is None:
        guess = {}
        for patch in patches:
            guess[patch] = patch.get_emission()

    # unlike Jacobi iteration, we update the residuals continuously during a step,
    # instead of saving the residuals for the next step's calculations
    len_patches = len(patches)
    converged = False
    while not converged:
        max_soln_change = blackRGBString()
        for i in xrange(len_patches):
            patch = patches[i]

            # compute form factors
            if cached_form_factors.has_key(patch):
                form_factors = cached_form_factors[patch]
            else:
                form_factors = formfactorengine.computeFormFactorsForPatch(patch)
                
            # compute K_ii (self form-factor)
            if form_factors.has_key(patch):
                F_ii = form_factors[patch]
            else:
                F_ii = formfactorengine.defaultFormFactor()
            K_ii = subtractRGBStrings(whiteRGBString(), multiplyRGBStringBy(patch.get_reflectivity(), F_ii))

            # B^(k+1) = E - K_j/K_ii*B_j^k
            B_k = guess[patch]
            guess[patch] = patch.get_emission()
            rho_i = patch.get_reflectivity()
            for patch_index in form_factors:
                if patch_index == i or patch_index >= len_patches:
                    continue
                
                # compute K_ji
                F_ij = form_factors[patch_index]
                K_ij = multiplyRGBStringBy(rho_i, -F_ij)
                
                guess[patch] = subtractRGBStrings(guess[patch], multiplyRGBStrings(guess[patches[patch_index]], divideRGBStrings(K_ij, K_ii)))

            # track change for convergence testing
            change = subtractRGBStrings(guess[patch], B_k)
            if isRGBStringGreaterThan(change, max_soln_change):
                max_soln_change = change

        # test for convergence
        if isRGBStringGreaterThan(convergence_threshold, max_soln_change):
            converged = True

        # notify after all the patches have been updated in this step
        if notifier is not None:
            notifier.notifyUpdatedGuess(guess, radiosityRelaxation_GaussSeidelIteration)
            notifier.notifyUpdatedConvergenceMetric(max_soln_change, radiosityRelaxation_GaussSeidelIteration)

    return guess

def radiosityRelaxation_SouthwellIteration(patches, formfactorengine, convergence_threshold="0.00001,0.00001,0.00001", cached_form_factors=None, formfactor_threshold=0.0, notifier=None):
    # unlike Jacobi or Gauss-Seidel, we use our own initial guess of 0
    # and initial residuals equal to the patch emissions
    guess = {}
    residuals = {}
    for patch in patches:
        guess[patch] = blackRGBString()
        residuals[patch] = patch.get_emission()

    # during each step, we choose the largest residual in the
    # guess, relax this residual, then update all the other
    # residuals
    converged = False
    notify_count = 0
    len_patches = len(patches)
    while not converged:
        # get the largest residual
        max_residual = blackRGBString()
        max_residual_patch = patches[0] # to ensure it's never None
        for patch in residuals:
            if isRGBStringGreaterThan(residuals[patch], max_residual):
                max_residual = residuals[patch]
                max_residual_patch = patch
        max_residual_patch_index = patches.index(max_residual_patch)

        # compute K_ii (self form-factor)
        if cached_form_factors.has_key(max_residual_patch):
            F_i = cached_form_factors[max_residual_patch]
            if F_i.has_key(max_residual_patch_index):
                F_ii = F_i[max_residual_patch_index]
            else:
                F_ii = formfactorengine.defaultFormFactor()
        else:
            F_ii = formfactorengine.computeFormFactorForPatchiToPatchj(max_residual_patch, max_residual_patch)
        K_ii = subtractRGBStrings(whiteRGBString(), multiplyRGBStringBy(max_residual_patch.get_reflectivity(), F_ii))
        
        # B^(k+1) = B^k + r^k/K_ii
        guess[max_residual_patch] = addRGBStrings(guess[max_residual_patch], divideRGBStrings(max_residual, K_ii))

        for patch in residuals:
            # r^(k+1) = r^k - max_residual*K_ji/K_ii

            if patch == max_residual_patch:
                K_ji = K_ii
            else:
                if cached_form_factors.has_key(patch):
                    F_j = cached_form_factors[patch]
                    if F_j.has_key(max_residual_patch_index):
                        F_ji = F_j[max_residual_patch_index]
                    else:
                        F_ji = formfactorengine.defaultFormFactor()
                else:
                    F_ji = formfactorengine.computeFormFactorForPatchiToPatchj(patch, max_residual_patch)
                K_ji = multiplyRGBStringBy(patch.get_reflectivity(), -F_ji)
                
            residuals[patch] = subtractRGBStrings(residuals[patch], multiplyRGBStrings(max_residual, divideRGBStrings(K_ji, K_ii)))

        if isRGBStringGreaterThan(convergence_threshold, max_residual):
            converged = True

        # notify after all the patches have been updated in this step
        # to compare accurately to Gauss-Seidel or Jacobi, we have to perform at least
        # as many updates as they do before notifying,
        # which is equal to the number of patches
        if notifier is not None:
            notify_count += 1
            if notify_count == len_patches:
                notifier.notifyUpdatedGuess(guess, radiosityRelaxation_SouthwellIteration)
                notifier.notifyUpdatedResidual(residuals, radiosityRelaxation_SouthwellIteration)
                notifier.notifyUpdatedConvergenceMetric(max_residual, radiosityRelaxation_SouthwellIteration)
                
                notify_count = 0
            
    return guess

def radiosityRelaxation_ProgressiveRefinement(patches, formfactorengine, convergence_threshold="0.00001,0.00001,0.00001", cached_form_factors=None, formfactor_threshold=0.0, display_function=None, notifier=None):
    # in progressive refinement, we use a guess and a delta guess, both
    # initialized to the patch's emission
    guess = {}
    delta_guess = {}
    for patch in patches:
        guess[patch] = patch.get_emission()
        delta_guess[patch] = patch.get_emission()

    # progressive refinement and southwell iteration are almost identical;
    # p.r. is equivalent to s.i. plus one full jacobi iteration (before display)
    converged = False
    notify_count = 0
    len_patches = len(patches)
    while not converged:
        # get the largest delta(B_i)*A_i
        largest_unshot_radiosity = blackRGBString()
        largest_unshot_radiosity_patch = patches[0] # to ensure it's never None
        for patch in delta_guess:
            area = triangleArea(patch.triangle)
            unshot_radiosity = multiplyRGBStringBy(delta_guess[patch], area)
            if isRGBStringGreaterThan(unshot_radiosity, largest_unshot_radiosity):
                largest_unshot_radiosity = unshot_radiosity
                largest_unshot_radiosity_patch = patch
        largest_unshot_radiosity_patch_index = patches.index(largest_unshot_radiosity_patch)

        max_soln_change = blackRGBString()
        for patch in patches:
            if cached_form_factors.has_key(patch):
                F_j = cached_form_factors[patch]
                if F_j.has_key(largest_unshot_radiosity_patch_index):
                    F_ji = F_j[largest_unshot_radiosity_patch_index]
                else:
                    F_ji = formfactorengine.defaultFormFactor()
            else:
                F_ji = formfactorengine.computeFormFactorForPatchiToPatchj(patch, largest_unshot_radiosity_patch)

            # del_rad = delta(B_i)*p_i*F_ji
            # delta(B_j) = delta(B_j) + del_rad
            # B_j = B_j + del_rad
            del_rad = multiplyRGBStrings(delta_guess[largest_unshot_radiosity_patch], multiplyRGBStringBy(patch.get_reflectivity(), F_ji))
            delta_guess[patch] = addRGBStrings(delta_guess[patch], del_rad)
            B_j = addRGBStrings(guess[patch], del_rad)

            change = subtractRGBStrings(B_j, guess[patch])
            if isRGBStringGreaterThan(change, max_soln_change):
                max_soln_change = change

            guess[patch] = B_j
            
        delta_guess[largest_unshot_radiosity_patch] = blackRGBString()

        if display_function is not None:
            display_function()

        if isRGBStringGreaterThan(convergence_threshold, max_soln_change):
            converged = True

        # notify after all the patches have been updated in this step
        # to compare accurately to Gauss-Seidel or Jacobi, we have to perform at least
        # as many updates as they do before notifying,
        # which is equal to the number of patches
        if notifier is not None:
            notify_count += 1
            if notify_count == len(patches):
                notifier.notifyUpdatedGuess(guess, radiosityRelaxation_ProgressiveRefinement)
                notifier.notifyUpdatedConvergenceMetric(largest_unshot_radiosity, radiosityRelaxation_ProgressiveRefinement)

                notify_count = 0

    return guess

def updatePatchesInGuess(guess):
    for patch in guess:
        patch.set_radiosity(guess[patch])

if __name__ == '__main__':
    from xmlscenereader import readScene
    from xmlscenewriter import writeScene

    class TestRadiosityRelaxation(unittest.TestCase):
        
        def setUp(self):
            self.infile = "./xml_test_files/test_scene_radiosity_relaxation.xml"
            self.reffile = "./xml_test_files/test_scene_rr_rrv_computed.xml"

            self.outfileJacobi = "./xml_test_files/test_scene_jacobi_iteration.xml"
            self.outfileGaussSeidel = "./xml_test_files/test_scene_gauss_seidel.xml"
            self.outfileSouthwell = "./xml_test_files/test_scene_southwell.xml"
            self.outfileProgressiveRefinement = "./xml_test_files/test_scene_progressive_refinement.xml"

            self.resolution = 256

            self.inscene = readScene(self.infile)
            self.patches = createPatches(self.inscene, applyMacros(self.inscene), 1000000.0) # to ensure no subdivision
            
            self.formfactorengine = HemicubeFormFactorEngine(patches=self.patches, resolution=self.resolution)
            # cache all the form factors for every patch
            self.cached_form_factors = {}
            for patch in self.patches:
                self.cached_form_factors[patch] = self.formfactorengine.computeFormFactorsForPatch(patch)

        def rr(self, radiosity_function, outfilename):
            guess = radiosity_function(patches=self.patches, formfactorenginer=self.formfactorenginer, cached_form_factors=self.cached_form_factors)
            updatePatchesInGuess(guess)
            computed_scene = convertPatchesToXmlScene(self.patches)
            writeScene(computed_scene, outfilename)

        def test_jacobi_iteration(self):
            self.rr(radiosityRelaxation_JacobiIteration, self.outfileJacobi)
            # TODO: ensure correct computation

        def test_gauss_seidel_iteration(self):
            self.rr(radiosityRelaxation_GaussSeidelIteration, self.outfileGaussSeidel)
            # TODO: ensure correct computation

        def test_southwell_iteration(self):
            self.rr(radiosityRelaxation_SouthwellIteration, self.outfileSouthwell)
            # TODO: ensure correct computation

        def test_progressive_refinement(self):
            self.rr(radiosityRelaxation_ProgressiveRefinement, self.outfileProgressiveRefinement)
            # TODO: ensure correct computation

    unittest.main()
